
/***************************************************************************//**
 * @file camera_array_depth.cpp
 * @brief Generates depth maps for image pairs and a combined depth map using the individual depth maps.
 *
 * @page "Camera Depth" Camera Depth
 *
 * @authors Brian Fehrman and Scott Logan
 * 
 * @date April 17 2013
 * 
 * @section file_section File Information
 *
 * @details Implementation file for the camera_array_depth functions
 *       which are used to generate depth maps for camera pairs as well as
 *       combine the individual depth maps in to a single depth map.
 * 
 * @section settings_section Required Settings
 * 
@verbatim
[ "base_dir" ] - name of folder where the matches.xml file is stored (string)
[ "output_dir" ] - name of folder where the depth map images should be saved (string)
[ "f_pixels" ] f_pixels - focal length in units of pixels (int)
[ "base_line" ] base_line - distance between horizontal/vertical camera pairs in units of mm (double)
[ "select_method" ] select_method - specifies how pixels are selected if there is a large disagreement among depth maps during combination (int)
[ "pixel_radius" ] pixel_radius - how large the drawn pixels will be (int)
@endverbatim
 * 
 * @section todo_bugs_modification_section Todo, Bugs, and Modifications
 *
 * @todo Improve combination algorithm
 ******************************************************************************/

//Includes
#include "camera_array_depth.h"
void generate_hybrid_image( Mat& depth, Mat& orig_image, Mat& hybrid_img, string output_dir );


/***************************************************************************//**
 * @author Brian Fehrman
 *
 * @par Description: Generates depth maps for a camera array. Using matching information
 * found for each pair of cameras and determines the distance between points across
 * the two images. A depth map is generated for each camera pair. These individual
 * depth maps are then combined in to a single depth map.
 *
 * @param[in] fs - FileNode pointing to the settings that are needed
 *
 ******************************************************************************/
void compute_array_depth( FileNode fs )
{
   string base_dir;
   string output_dir;
   float f_pixels;
   float base_line;
   int select_method;
   int pixel_radius;
   
   fs[ "base_dir" ] >> base_dir;
   fs[ "output_dir" ] >> output_dir;
   fs[ "f_pixels" ] >> f_pixels;
   fs[ "base_line" ] >> base_line;
   fs[ "select_method" ] >> select_method;
   fs[ "pixel_radius" ] >> pixel_radius;
   
   Mat average_depth;
   FileStorage match_file;
   vector< Mat > matches;
   vector< Point > search_direction;
   vector< Mat > diff_mats;
   int max_search;

   base_dir += "/matches.xml";
   
   mkdir( output_dir.c_str(), S_IRWXU|S_IRGRP|S_IXGRP );   
   output_dir += "/";
   
   match_file.open( base_dir, FileStorage::READ );
   
   match_file[ "matches" ] >> matches;
   match_file[ "search_direction" ] >> search_direction;
   
   if( search_direction[ 0 ].x > search_direction[ 0 ].y )
   {
      max_search = search_direction[ 0 ].x;
   }
   else
   {
      max_search = search_direction[ 0 ].y; 
   }
   
   for( unsigned int match = 0; match < matches.size(); match++ )
   {
      Mat diff_mat;
      bool x_change = false, y_change = false;
      
      if( search_direction[ match ].x == search_direction[ match ].y )
      {
         x_change = true;
         y_change = true;
      }
      else if( search_direction[ match ].x > search_direction[ match ].y )
      {
         x_change = true;
      }
      else
      {
         y_change = true;
      }
      
      compute_diff_mat( matches[ match ], diff_mat, x_change, y_change );
      
      diff_mats.push_back( diff_mat );
   }
   
   combine_depth_maps( matches, diff_mats, average_depth, select_method );
   generate_color_maps( diff_mats, average_depth, base_line, f_pixels, max_search, output_dir, pixel_radius );
}

/***************************************************************************//**
 * @author Brian Fehrman
 *
 * @par Description: Computes the disparity between matching points. It will only
 * consider the distance along the main camera offset. For diagonal pairs, it
 * will account for both directions and then scale the result by 1/sqrt(2).
 *
 * @param[in] matches - Contains the matching information for a camera pair
 * @param[out] diff_mat - stores the disparity between the images
 * @param[in] x_change - if the main camera offset was along the x direction
 * @param[in] y_change - if the main camera offset was along the y direction
 *
 ******************************************************************************/
void compute_diff_mat( Mat matches, Mat& diff_mat, bool x_change, bool y_change )
{
   int nrows = matches.rows;
   int ncols = matches.cols;
   double x_diff, y_diff;
   
   diff_mat = Mat::zeros( nrows, ncols, CV_8UC1 );
   
   for( int row = 0; row < nrows; row++ )
   {
      for( int col = 0; col < ncols; col++ )
      {
         if( matches.at< Vec3f >( row, col )[ 2 ] > 0 )
         {
            x_diff = abs( matches.at< Vec3f >( row, col )[ 0 ] - col ); 
            y_diff = abs( matches.at< Vec3f >( row, col )[ 1 ] - row );
            
            //Divide by sqrt(2) to account for longer diagonal baseline
            if( x_change && y_change )
            {
               diff_mat.at< unsigned char >( row, col ) = sqrt( x_diff * x_diff + y_diff * y_diff ) / sqrt( 2 );
            }
            else if( y_change )
            {
               diff_mat.at< unsigned char >( row, col ) = y_diff;
            }
            else
            {
               diff_mat.at< unsigned char >( row, col ) = x_diff;
            }
         }
      }
   }
}

/***************************************************************************//**
 * @author Brian Fehrman
 *
 * @par Description: Computes the disparity between matching points. It will only
 * consider the distance along the main camera offset. For diagonal pairs, it
 * will account for both directions and then scale the result by 1/sqrt(2).
 *
 * @param[in] matches - Contains the matching information for all camera pairs, the confidence information from this is needed
 * @param[in] diff_mats - The disparities found for all camera pairs
 * @param[out] average_depth - The combined depth map
 * @param[in] select_method - determines what to do when there is a large disagreement among the depth maps for a given pixel
 *
 ******************************************************************************/
void combine_depth_maps( vector< Mat > matches, vector< Mat > diff_mats, Mat& average_depth, int select_method )
{
   int nrows = diff_mats[ 0 ].rows;
   int ncols = diff_mats[ 0 ].cols;
   
   average_depth = Mat::zeros( nrows, ncols, CV_8UC1 );
   
   for( int row = 0; row < nrows; row++ )
   {
      for( int col = 0; col < ncols; col++ )
      {
         float avg = 0.0;
         float std = 0.0;
         vector<float> vals_used;
         vector<float> conf_used;
         float max_conf = -1;
         int max_conf_idx = 0;
         float total_conf = 0.0;
         
         for( unsigned int idx = 0; idx < diff_mats.size(); idx++ )
         {
            //matches[ idx ].at<Vec3f>( row, col )[ 2 ]  is the confidence of the match
            if( matches[ idx ].at<Vec3f>( row, col )[ 2 ] > 0 )
            {
               vals_used.push_back( diff_mats[ idx ].at< unsigned char>( row, col )  );
               total_conf += matches[ idx ].at<Vec3f>( row, col )[ 2 ];
               conf_used.push_back( matches[ idx ].at<Vec3f>( row, col )[ 2 ] );
            
               if( matches[ idx ].at<Vec3f>( row, col )[ 2 ] > max_conf )
               {
                  max_conf = matches[ idx ].at<Vec3f>( row, col )[ 2 ];
                  max_conf_idx = vals_used.size() - 1;
               }
            }
         }
         
         for( unsigned int idx = 0; idx < vals_used.size(); idx++ )
         {
            avg += vals_used[ idx ];
         }
         
         avg /= vals_used.size() + 0.5;
         
         for( unsigned int idx = 0; idx < vals_used.size(); idx++ )
         {
            std += ( vals_used[ idx ] - avg ) * ( vals_used[ idx ] - avg );
         }
         
         std /= vals_used.size();
         
         std = sqrt( std );
         
         if( std > 12  )
         {
            //What to do if large disagreement among distance measurements
            switch( select_method )
            {
               //Most confident value
               case 0:
                  average_depth.at< unsigned char >( row, col ) = vals_used[ max_conf_idx ];
                  break;
                  
               //Closest value
               case 1:
                  average_depth.at< unsigned char >( row, col ) = *max_element( vals_used.begin(), vals_used.end() );
                  break;
            }
         }
         else
         {
            avg = 0.0;
            for( unsigned int idx = 0; idx < vals_used.size(); idx++ )
            {
               avg += vals_used[ idx ] * conf_used[ idx ] / total_conf ;
            }
            
            average_depth.at< unsigned char >( row, col ) = avg;
         }
      }
   }
   
   //medianBlur( average_depth, average_depth, 7);
}

/***************************************************************************//**
 * @author Brian Fehrman and Scott Logan
 *
 * @par Description: Generates color depth maps based on the disparities found among camera pairs.
 * A scale with the estimated real world distance in mm is shown to the right of each depth map. Unknown
 * depths are shown as a gray color.
 * 
 * @param[in] diff_mats - The disparities found for all camera pairs
 * @param[in] average_depth - The combined depth map
 * @param[in] base_line - distance in mm between camera pairs (used to estimate true depth in mm)
 * @param[in] f_pixels - approximate focal length of cameras in units of pixels
 * @param[in] max_search - the largest possible pixel displacement between images
 * @param[in] output_dir - where the colored depth maps should be saved
 * @param[in] pixel_radius - determines the size of the pixels that are drawn
 *
 ******************************************************************************/
void generate_color_maps( vector< Mat > diff_mats, Mat& average_depth, float base_line, float f_pixels, int max_search, string output_dir, const int pixel_radius )
{
   int nrows =  average_depth.rows;
   int ncols = average_depth.cols;
   
   diff_mats.push_back( average_depth );
   
   for( unsigned int idx = 0; idx < diff_mats.size(); idx++ )
   {
      Mat diff_color( nrows, ncols + 50, CV_8UC3 );
      stringstream ss;
      ss << idx;

      for( int row = 0; row < nrows; row++ )
      {
          for( int col = 0; col < ncols; col++ )
          {
              diff_color.at< Vec3b >( row, col )[ 0 ] = 0;
              diff_color.at< Vec3b >( row, col )[ 1 ] = 0;
              diff_color.at< Vec3b >( row, col )[ 2 ] = 128;
          }
      }
      
      for( int row = 0; row < nrows; row++ )
      {
         for( int col = 0; col < ncols; col++ )
         {
            for(int i = (row - pixel_radius < 0 ) ? -pixel_radius - (row - pixel_radius): -pixel_radius; i <= pixel_radius && row + i < nrows; i++)
            {
                for(int j = (col - pixel_radius < 0 ) ? -pixel_radius - (col - pixel_radius): -pixel_radius; j <= pixel_radius && col + j < ncols; j++ )
                {
                    if( diff_mats[ idx ].at< unsigned char >( row, col ) * ( 150.0 / max_search ) > diff_color.at< Vec3b >( row + i, col + j )[ 0 ] )
                    {
                        if( diff_mats[ idx ].at< unsigned char >( row, col ) > 0 )
                        {
                           diff_color.at< Vec3b >( row + i, col + j )[ 0 ] = diff_mats[ idx ].at< unsigned char >( row, col ) * ( 150.0 / max_search );
                           diff_color.at< Vec3b >( row + i, col + j )[ 1 ] = 255;
                        }
                     }
                }
            }
         }
      }
      
      for( int row =  0; row < nrows; row++ )
      {
         for( int col = ncols; col < ncols + 50; col++ )
         {
            diff_color.at< Vec3b >( row, col )[ 0 ] =  ( float ) row / nrows * 150;
            diff_color.at< Vec3b >( row, col )[ 1 ] = 255;
            diff_color.at< Vec3b >( row, col )[ 2 ] = 255;
         }
      }
      
      for( int row = 10; row < nrows; row += 10 )
      {
            stringstream distance;
                        
            distance << int( ( base_line * f_pixels ) / ( ( float ) row / nrows * max_search ) );
            
            putText( diff_color, distance.str(), Point( ncols, row ), CV_FONT_HERSHEY_SIMPLEX, .35, Scalar( 0, 0 ,0 ) , 1 );
      }

      cvtColor( diff_color, diff_color, CV_HSV2RGB );
      imwrite( output_dir + ss.str() + ".jpeg", diff_color );
      
      if( idx == diff_mats.size() - 1 )
      {
            Mat orig_img = imread( "orig_1_1.jpg" );
            Mat hy_img;
   
            generate_hybrid_image( diff_color, orig_img, hy_img, "depth/" );
      }
   }
}

void generate_hybrid_image( Mat& depth, Mat& orig_image, Mat& hybrid_img, string output_dir )
{
   int nrows = depth.rows;
   int ncols = depth.cols;
   
   hybrid_img = Mat::zeros( nrows, ncols, CV_8UC3 );
   
   cvtColor( orig_image, orig_image, CV_RGB2HSV );
   
   cvtColor(  depth, depth, CV_RGB2HSV );
   
   for( int row = 0; row < nrows; row++ )
   {
      for( int col = 0; col < ncols; col++ )
      {      
         if( depth.at< Vec3b >( row, col )[ 1 ] == 0 )
         {      
            hybrid_img.at< Vec3b >( row, col ) = orig_image.at< Vec3b >( row, col );
         }
         else
         {
            hybrid_img.at< Vec3b >( row, col ) = depth.at< Vec3b >( row, col );
         }
      }
   }
   
   cvtColor( hybrid_img, hybrid_img, CV_HSV2RGB );
   imwrite( output_dir + "hybrid_img.jpeg", hybrid_img );
}
